package com.ai.langchain4j.controller;

import java.time.LocalDate;

import com.ai.langchain4j.config.Langchain4JConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.service.TokenStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @Author: suren@517na.com
 * @CreateTime: 2025-05-25
 * @Description: springboot整合Langchain4j Controlelr
 */

@RestController
@RequestMapping("/v2/ai")
public class SpringLangchain4jAiController {

    @Autowired
    QwenChatModel qwenChatModel;

    @Autowired
    QwenStreamingChatModel qwenStreamingChatModel;

    @Autowired
    OllamaChatModel langchain4jOllamaChatModel;

    @Autowired
    Langchain4JConfig.Assistant assistant;

    @Autowired
    Langchain4JConfig.AssistantUnique assistantUnique;


    @GetMapping("/ollama/chat")
    public String ollamaChat(@RequestParam(defaultValue = "你是谁") String message) {
        return langchain4jOllamaChatModel.chat(message);
    }


    @GetMapping("/chat")
    public String chat(@RequestParam(defaultValue = "你是谁") String message) {
        return qwenChatModel.chat(message);
    }


    @GetMapping(value = "/stream/chat", produces = "text/stream;charset=UTF-8")
    public Flux<String> chatStream(@RequestParam(defaultValue = "你是谁") String message) {
        Flux<String> stringFlux = Flux.create(sink -> {
            qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    sink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    sink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    sink.error(throwable);
                }
            });
        });
        return stringFlux;
    }


    /**
     * 记忆聊天
     * @param message
     * @return
     */
    @GetMapping("/memory/chat")
    public String memoryChat(@RequestParam(defaultValue = "你是谁") String message) {
        return assistant.chat(message);
    }


    /**
     * 记忆聊天-流式输出
     * @param message
     * @return
     */
    @GetMapping(value = "/stream/memory/chat", produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam(defaultValue = "你是谁") String message) {
        TokenStream stream = assistant.stream(message);
        return Flux.create(sink -> {
            stream.onPartialResponse( s-> sink.next(s))
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();
        });
    }

    /**
     * 记忆聊天 - 区分用户
     * @param message
     * @return
     */
    @GetMapping("/memory_id/chat")
    public String memoryIdChat(@RequestParam(defaultValue = "你是谁") String message, @RequestParam Integer userId) {
        return assistantUnique.chat(userId, message);
    }

    @GetMapping(value = "/stream/memory_id/chat", produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryIdStreamChat(@RequestParam(defaultValue = "你是谁") String message, @RequestParam Integer userId) {
        TokenStream stream = assistantUnique.stream(userId, message, LocalDate.now().toString());
        return Flux.create(sink -> {
            stream.onPartialResponse( s-> sink.next(s))
                    .onCompleteResponse(c -> sink.complete())
                    .onError(sink::error)
                    .start();
        });
    }
}
